home *** CD-ROM | disk | FTP | other *** search
/ MacWorld 2005 March / Macworld CD March 2005 - Marathon Trilogy.iso / Shareware World / iPod / iPodderX.sit / iPodderX / iPodderX.app / Contents / Resources / StorageWrapper.py < prev    next >
Encoding:
Python Source  |  2005-01-07  |  16.6 KB  |  470 lines
  1. # Written by Bram Cohen
  2. # see LICENSE.txt for license information
  3.  
  4. from sha import sha
  5. from threading import Event
  6. from bitfield import Bitfield
  7.  
  8. def dummy_status(fractionDone = None, activity = None):
  9.     pass
  10.  
  11. def dummy_data_flunked(size):
  12.     pass
  13.  
  14. class StorageWrapper:
  15.     def __init__(self, storage, request_size, hashes, 
  16.             piece_size, finished, failed, 
  17.             statusfunc = dummy_status, flag = Event(), check_hashes = True,
  18.             data_flunked = dummy_data_flunked):
  19.         self.storage = storage
  20.         self.request_size = request_size
  21.         self.hashes = hashes
  22.         self.piece_size = piece_size
  23.         self.data_flunked = data_flunked
  24.         self.total_length = storage.get_total_length()
  25.         self.amount_left = self.total_length
  26.         if self.total_length <= piece_size * (len(hashes) - 1):
  27.             raise ValueError, 'bad data from tracker - total too small'
  28.         if self.total_length > piece_size * len(hashes):
  29.             raise ValueError, 'bad data from tracker - total too big'
  30.         self.finished = finished
  31.         self.failed = failed
  32.         self.numactive = [0] * len(hashes)
  33.         self.inactive_requests = [1] * len(hashes)
  34.         self.amount_inactive = self.total_length
  35.         self.endgame = False
  36.         self.have = Bitfield(len(hashes))
  37.         self.waschecked = [check_hashes] * len(hashes)
  38.         self.places = {}
  39.         self.holes = []
  40.         if len(hashes) == 0:
  41.             finished()
  42.             return
  43.         targets = {}
  44.         total = len(hashes)
  45.         for i in xrange(len(hashes)):
  46.             if not self._waspre(i):
  47.                 targets.setdefault(hashes[i], []).append(i)
  48.                 total -= 1
  49.         numchecked = 0.0
  50.         if total and check_hashes:
  51.             statusfunc({"activity" : 'checking existing file', 
  52.                 "fractionDone" : 0})
  53.         def markgot(piece, pos, self = self, check_hashes = check_hashes):
  54.             self.places[piece] = pos
  55.             self.have[piece] = True
  56.             self.amount_left -= self._piecelen(piece)
  57.             self.amount_inactive -= self._piecelen(piece)
  58.             self.inactive_requests[piece] = None
  59.             self.waschecked[piece] = check_hashes
  60.         lastlen = self._piecelen(len(hashes) - 1)
  61.         for i in xrange(len(hashes)):
  62.             if not self._waspre(i):
  63.                 self.holes.append(i)
  64.             elif not check_hashes:
  65.                 markgot(i, i)
  66.             else:
  67.                 sh = sha(self.storage.read(piece_size * i, lastlen))
  68.                 sp = sh.digest()
  69.                 sh.update(self.storage.read(piece_size * i + lastlen, self._piecelen(i) - lastlen))
  70.                 s = sh.digest()
  71.                 if s == hashes[i]:
  72.                     markgot(i, i)
  73.                 elif targets.get(s) and self._piecelen(i) == self._piecelen(targets[s][-1]):
  74.                     markgot(targets[s].pop(), i)
  75.                 elif not self.have[len(hashes) - 1] and sp == hashes[-1] and (i == len(hashes) - 1 or not self._waspre(len(hashes) - 1)):
  76.                     markgot(len(hashes) - 1, i)
  77.                 else:
  78.                     self.places[i] = i
  79.                 if flag.isSet():
  80.                     return
  81.                 numchecked += 1
  82.                 statusfunc({'fractionDone': 1 - float(self.amount_left) / self.total_length})
  83.         if self.amount_left == 0:
  84.             finished()
  85.  
  86.     def _waspre(self, piece):
  87.         return self.storage.was_preallocated(piece * self.piece_size, self._piecelen(piece))
  88.  
  89.     def _piecelen(self, piece):
  90.         if piece < len(self.hashes) - 1:
  91.             return self.piece_size
  92.         else:
  93.             return self.total_length - piece * self.piece_size
  94.  
  95.     def get_amount_left(self):
  96.         return self.amount_left
  97.  
  98.     def do_I_have_anything(self):
  99.         return self.amount_left < self.total_length
  100.  
  101.     def _make_inactive(self, index):
  102.         length = min(self.piece_size, self.total_length - self.piece_size * index)
  103.         l = []
  104.         x = 0
  105.         while x + self.request_size < length:
  106.             l.append((x, self.request_size))
  107.             x += self.request_size
  108.         l.append((x, length - x))
  109.         self.inactive_requests[index] = l
  110.  
  111.     def is_endgame(self):
  112.         return self.endgame
  113.  
  114.     def get_have_list(self):
  115.         return self.have.tostring()
  116.  
  117.     def do_I_have(self, index):
  118.         return self.have[index]
  119.  
  120.     def do_I_have_requests(self, index):
  121.         return not not self.inactive_requests[index]
  122.  
  123.     def new_request(self, index):
  124.         # returns (begin, length)
  125.         if self.inactive_requests[index] == 1:
  126.             self._make_inactive(index)
  127.         self.numactive[index] += 1
  128.         rs = self.inactive_requests[index]
  129.         r = min(rs)
  130.         rs.remove(r)
  131.         self.amount_inactive -= r[1]
  132.         if self.amount_inactive == 0:
  133.             self.endgame = True
  134.         return r
  135.  
  136.     def piece_came_in(self, index, begin, piece):
  137.         try:
  138.             return self._piece_came_in(index, begin, piece)
  139.         except IOError, e:
  140.             self.failed('IO Error ' + str(e))
  141.             return True
  142.  
  143.     def _piece_came_in(self, index, begin, piece):
  144.         if not self.places.has_key(index):
  145.             n = self.holes.pop(0)
  146.             if self.places.has_key(n):
  147.                 oldpos = self.places[n]
  148.                 old = self.storage.read(self.piece_size * oldpos, self._piecelen(n))
  149.                 if self.have[n] and sha(old).digest() != self.hashes[n]:
  150.                     self.failed('data corrupted on disk - maybe you have two copies running?')
  151.                     return True
  152.                 self.storage.write(self.piece_size * n, old)
  153.                 self.places[n] = n
  154.                 if index == oldpos or index in self.holes:
  155.                     self.places[index] = oldpos
  156.                 else:
  157.                     for p, v in self.places.items():
  158.                         if v == index:
  159.                             break
  160.                     self.places[index] = index
  161.                     self.places[p] = oldpos
  162.                     old = self.storage.read(self.piece_size * index, self.piece_size)
  163.                     self.storage.write(self.piece_size * oldpos, old)
  164.             elif index in self.holes or index == n:
  165.                 if not self._waspre(n):
  166.                     self.storage.write(self.piece_size * n, self._piecelen(n) * chr(0xFF))
  167.                 self.places[index] = n
  168.             else:
  169.                 for p, v in self.places.items():
  170.                     if v == index:
  171.                         break
  172.                 self.places[index] = index
  173.                 self.places[p] = n
  174.                 old = self.storage.read(self.piece_size * index, self._piecelen(n))
  175.                 self.storage.write(self.piece_size * n, old)
  176.         self.storage.write(self.places[index] * self.piece_size + begin, piece)
  177.         self.numactive[index] -= 1
  178.         if not self.inactive_requests[index] and not self.numactive[index]:
  179.             if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() == self.hashes[index]:
  180.                 self.have[index] = True
  181.                 self.inactive_requests[index] = None
  182.                 self.waschecked[index] = True
  183.                 self.amount_left -= self._piecelen(index)
  184.                 if self.amount_left == 0:
  185.                     self.finished()
  186.             else:
  187.                 self.data_flunked(self._piecelen(index))
  188.                 self.inactive_requests[index] = 1
  189.                 self.amount_inactive += self._piecelen(index)
  190.                 return False
  191.         return True
  192.  
  193.     def request_lost(self, index, begin, length):
  194.         self.inactive_requests[index].append((begin, length))
  195.         self.amount_inactive += length
  196.         self.numactive[index] -= 1
  197.  
  198.     def get_piece(self, index, begin, length):
  199.         try:
  200.             return self._get_piece(index, begin, length)
  201.         except IOError, e:
  202.             self.failed('IO Error ' + str(e))
  203.             return None
  204.  
  205.     def _get_piece(self, index, begin, length):
  206.         if not self.have[index]:
  207.             return None
  208.         if not self.waschecked[index]:
  209.             if sha(self.storage.read(self.piece_size * self.places[index], self._piecelen(index))).digest() != self.hashes[index]:
  210.                 self.failed('told file complete on start-up, but piece failed hash check')
  211.                 return None
  212.             self.waschecked[index] = True
  213.         if begin + length > self._piecelen(index):
  214.             return None
  215.         return self.storage.read(self.piece_size * self.places[index] + begin, length)
  216.  
  217. class DummyStorage:
  218.     def __init__(self, total, pre = False, ranges = []):
  219.         self.pre = pre
  220.         self.ranges = ranges
  221.         self.s = chr(0xFF) * total
  222.         self.done = False
  223.  
  224.     def was_preexisting(self):
  225.         return self.pre
  226.  
  227.     def was_preallocated(self, begin, length):
  228.         for b, l in self.ranges:
  229.             if begin >= b and begin + length <= b + l:
  230.                 return True
  231.         return False
  232.  
  233.     def get_total_length(self):
  234.         return len(self.s)
  235.  
  236.     def read(self, begin, length):
  237.         return self.s[begin:begin + length]
  238.  
  239.     def write(self, begin, piece):
  240.         self.s = self.s[:begin] + piece + self.s[begin + len(piece):]
  241.  
  242.     def finished(self):
  243.         self.done = True
  244.  
  245. def test_basic():
  246.     ds = DummyStorage(3)
  247.     sw = StorageWrapper(ds, 2, [sha('abc').digest()], 4, ds.finished, None)
  248.     assert sw.get_amount_left() == 3
  249.     assert not sw.do_I_have_anything()
  250.     assert sw.get_have_list() == chr(0)
  251.     assert sw.do_I_have_requests(0)
  252.     x = []
  253.     x.append(sw.new_request(0))
  254.     assert sw.do_I_have_requests(0)
  255.     x.append(sw.new_request(0))
  256.     assert not sw.do_I_have_requests(0)
  257.     x.sort()
  258.     assert x == [(0, 2), (2, 1)]
  259.     sw.request_lost(0, 2, 1)
  260.     del x[-1]
  261.     assert sw.do_I_have_requests(0)
  262.     x.append(sw.new_request(0))
  263.     assert x == [(0, 2), (2, 1)]
  264.     assert not sw.do_I_have_requests(0)
  265.     sw.piece_came_in(0, 0, 'ab')
  266.     assert not sw.do_I_have_requests(0)
  267.     assert sw.get_amount_left() == 3
  268.     assert not sw.do_I_have_anything()
  269.     assert sw.get_have_list() == chr(0)
  270.     assert not ds.done
  271.     sw.piece_came_in(0, 2, 'c')
  272.     assert not sw.do_I_have_requests(0)
  273.     assert sw.get_amount_left() == 0
  274.     assert sw.do_I_have_anything()
  275.     assert sw.get_have_list() == chr(0x80)
  276.     assert sw.get_piece(0, 0, 3) == 'abc'
  277.     assert sw.get_piece(0, 1, 2) == 'bc'
  278.     assert sw.get_piece(0, 0, 2) == 'ab'
  279.     assert sw.get_piece(0, 1, 1) == 'b'
  280.     assert ds.done
  281.  
  282. def test_two_pieces():
  283.     ds = DummyStorage(4)
  284.     sw = StorageWrapper(ds, 3, [sha('abc').digest(),
  285.         sha('d').digest()], 3, ds.finished, None)
  286.     assert sw.get_amount_left() == 4
  287.     assert not sw.do_I_have_anything()
  288.     assert sw.get_have_list() == chr(0)
  289.     assert sw.do_I_have_requests(0)
  290.     assert sw.do_I_have_requests(1)
  291.  
  292.     assert sw.new_request(0) == (0, 3)
  293.     assert sw.get_amount_left() == 4
  294.     assert not sw.do_I_have_anything()
  295.     assert sw.get_have_list() == chr(0)
  296.     assert not sw.do_I_have_requests(0)
  297.     assert sw.do_I_have_requests(1)
  298.  
  299.     assert sw.new_request(1) == (0, 1)
  300.     assert sw.get_amount_left() == 4
  301.     assert not sw.do_I_have_anything()
  302.     assert sw.get_have_list() == chr(0)
  303.     assert not sw.do_I_have_requests(0)
  304.     assert not sw.do_I_have_requests(1)
  305.  
  306.     sw.piece_came_in(0, 0, 'abc')
  307.     assert sw.get_amount_left() == 1
  308.     assert sw.do_I_have_anything()
  309.     assert sw.get_have_list() == chr(0x80)
  310.     assert not sw.do_I_have_requests(0)
  311.     assert not sw.do_I_have_requests(1)
  312.     assert sw.get_piece(0, 0, 3) == 'abc'
  313.     assert not ds.done
  314.  
  315.     sw.piece_came_in(1, 0, 'd')
  316.     assert ds.done
  317.     assert sw.get_amount_left() == 0
  318.     assert sw.do_I_have_anything()
  319.     assert sw.get_have_list() == chr(0xC0)
  320.     assert not sw.do_I_have_requests(0)
  321.     assert not sw.do_I_have_requests(1)
  322.     assert sw.get_piece(1, 0, 1) == 'd'
  323.  
  324. def test_hash_fail():
  325.     ds = DummyStorage(4)
  326.     sw = StorageWrapper(ds, 4, [sha('abcd').digest()], 4, ds.finished, None)
  327.     assert sw.get_amount_left() == 4
  328.     assert not sw.do_I_have_anything()
  329.     assert sw.get_have_list() == chr(0)
  330.     assert sw.do_I_have_requests(0)
  331.  
  332.     assert sw.new_request(0) == (0, 4)
  333.     sw.piece_came_in(0, 0, 'abcx')
  334.     assert sw.get_amount_left() == 4
  335.     assert not sw.do_I_have_anything()
  336.     assert sw.get_have_list() == chr(0)
  337.     assert sw.do_I_have_requests(0)
  338.  
  339.     assert sw.new_request(0) == (0, 4)
  340.     assert not ds.done
  341.     sw.piece_came_in(0, 0, 'abcd')
  342.     assert ds.done
  343.     assert sw.get_amount_left() == 0
  344.     assert sw.do_I_have_anything()
  345.     assert sw.get_have_list() == chr(0x80)
  346.     assert not sw.do_I_have_requests(0)
  347.  
  348. def test_lazy_hashing():
  349.     ds = DummyStorage(4, ranges = [(0, 4)])
  350.     flag = Event()
  351.     sw = StorageWrapper(ds, 4, [sha('abcd').digest()], 4, ds.finished, lambda x, flag = flag: flag.set(), check_hashes = False)
  352.     assert sw.get_piece(0, 0, 2) is None
  353.     assert flag.isSet()
  354.  
  355. def test_lazy_hashing_pass():
  356.     ds = DummyStorage(4)
  357.     flag = Event()
  358.     sw = StorageWrapper(ds, 4, [sha(chr(0xFF) * 4).digest()], 4, ds.finished, lambda x, flag = flag: flag.set(), check_hashes = False)
  359.     assert sw.get_piece(0, 0, 2) is None
  360.     assert not flag.isSet()
  361.  
  362. def test_preexisting():
  363.     ds = DummyStorage(4, True, [(0, 4)])
  364.     sw = StorageWrapper(ds, 2, [sha(chr(0xFF) * 2).digest(), 
  365.         sha('ab').digest()], 2, ds.finished, None)
  366.     assert sw.get_amount_left() == 2
  367.     assert sw.do_I_have_anything()
  368.     assert sw.get_have_list() == chr(0x80)
  369.     assert not sw.do_I_have_requests(0)
  370.     assert sw.do_I_have_requests(1)
  371.     assert sw.new_request(1) == (0, 2)
  372.     assert not ds.done
  373.     sw.piece_came_in(1, 0, 'ab')
  374.     assert ds.done
  375.     assert sw.get_amount_left() == 0
  376.     assert sw.do_I_have_anything()
  377.     assert sw.get_have_list() == chr(0xC0)
  378.     assert not sw.do_I_have_requests(0)
  379.     assert not sw.do_I_have_requests(1)
  380.  
  381. def test_total_too_short():
  382.     ds = DummyStorage(4)
  383.     try:
  384.         StorageWrapper(ds, 4, [sha(chr(0xff) * 4).digest(),
  385.             sha(chr(0xFF) * 4).digest()], 4, ds.finished, None)
  386.         raise 'fail'
  387.     except ValueError:
  388.         pass
  389.  
  390. def test_total_too_big():
  391.     ds = DummyStorage(9)
  392.     try:
  393.         sw = StorageWrapper(ds, 4, [sha('qqqq').digest(),
  394.             sha(chr(0xFF) * 4).digest()], 4, ds.finished, None)
  395.         raise 'fail'
  396.     except ValueError:
  397.         pass
  398.  
  399. def test_end_above_total_length():
  400.     ds = DummyStorage(3, True)
  401.     sw = StorageWrapper(ds, 4, [sha('qqq').digest()], 4, ds.finished, None)
  402.     assert sw.get_piece(0, 0, 4) == None
  403.  
  404. def test_end_past_piece_end():
  405.     ds = DummyStorage(4, True, ranges = [(0, 4)])
  406.     sw = StorageWrapper(ds, 4, [sha(chr(0xFF) * 2).digest(), 
  407.         sha(chr(0xFF) * 2).digest()], 2, ds.finished, None)
  408.     assert ds.done
  409.     assert sw.get_piece(0, 0, 3) == None
  410.  
  411. from random import shuffle
  412.  
  413. def test_alloc_random():
  414.     ds = DummyStorage(101)
  415.     sw = StorageWrapper(ds, 1, [sha(chr(i)).digest() for i in xrange(101)], 1, ds.finished, None)
  416.     for i in xrange(100):
  417.         assert sw.new_request(i) == (0, 1)
  418.     r = range(100)
  419.     shuffle(r)
  420.     for i in r:
  421.         sw.piece_came_in(i, 0, chr(i))
  422.     for i in xrange(100):
  423.         assert sw.get_piece(i, 0, 1) == chr(i)
  424.     assert ds.s[:100] == ''.join([chr(i) for i in xrange(100)])
  425.  
  426. def test_alloc_resume():
  427.     ds = DummyStorage(101)
  428.     sw = StorageWrapper(ds, 1, [sha(chr(i)).digest() for i in xrange(101)], 1, ds.finished, None)
  429.     for i in xrange(100):
  430.         assert sw.new_request(i) == (0, 1)
  431.     r = range(100)
  432.     shuffle(r)
  433.     for i in r[:50]:
  434.         sw.piece_came_in(i, 0, chr(i))
  435.     assert ds.s[50:] == chr(0xFF) * 51
  436.     ds.ranges = [(0, 50)]
  437.     sw = StorageWrapper(ds, 1, [sha(chr(i)).digest() for i in xrange(101)], 1, ds.finished, None)
  438.     for i in r[50:]:
  439.         sw.piece_came_in(i, 0, chr(i))
  440.     assert ds.s[:100] == ''.join([chr(i) for i in xrange(100)])
  441.  
  442. def test_last_piece_pre():
  443.     ds = DummyStorage(3, ranges = [(2, 1)])
  444.     ds.s = chr(0xFF) + chr(0xFF) + 'c'
  445.     sw = StorageWrapper(ds, 2, [sha('ab').digest(), sha('c').digest()], 2, ds.finished, None)
  446.     assert not sw.do_I_have_requests(1)
  447.     assert sw.do_I_have_requests(0)
  448.  
  449. def test_not_last_pre():
  450.     ds = DummyStorage(3, ranges = [(1, 1)])
  451.     ds.s = chr(0xFF) + 'a' + chr(0xFF)
  452.     sw = StorageWrapper(ds, 1, [sha('a').digest()] * 3, 1, ds.finished, None)
  453.     assert not sw.do_I_have_requests(1)
  454.     assert sw.do_I_have_requests(0)
  455.     assert sw.do_I_have_requests(2)
  456.  
  457. def test_last_piece_not_pre():
  458.     ds = DummyStorage(51, ranges = [(50, 1)])
  459.     sw = StorageWrapper(ds, 2, [sha('aa').digest()] * 25 + [sha('b').digest()], 2, ds.finished, None)
  460.     for i in xrange(25):
  461.         assert sw.new_request(i) == (0, 2)
  462.     assert sw.new_request(25) == (0, 1)
  463.     sw.piece_came_in(25, 0, 'b')
  464.     r = range(25)
  465.     shuffle(r)
  466.     for i in r:
  467.         sw.piece_came_in(i, 0, 'aa')
  468.     assert ds.done
  469.     assert ds.s == 'a' * 50 + 'b'
  470.